#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <stdint.h>
#include <inttypes.h>
#include <endian.h>
#include <byteswap.h>
#include <getopt.h>

#include <sys/time.h>
#include <arpa/inet.h>
#include <infiniband/verbs.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <netdb.h>

#define MAX_POLL_CQ_TIMEOUT 2000
#define SMALL_MSG_SIZE 32     // 16 bytes for small messages
#define MID_MSG_SIZE (512 * 1024)
#define LARGE_MSG_SIZE (32 * 1024 * 1024) // 16 MB for large messages
#define MAX_QPS 2001
#define MSG "SEND operation "
#define RDMAMSGR "RDMA read operation "
#define RDMAMSGW "RDMA write operation"
#define MSG_SIZE (strlen(MSG) + 1)
#if __BYTE_ORDER == __LITTLE_ENDIAN
static inline uint64_t htonll(uint64_t x) { return bswap_64(x); }
static inline uint64_t ntohll(uint64_t x) { return bswap_64(x); }
#elif __BYTE_ORDER == __BIG_ENDIAN
static inline uint64_t htonll(uint64_t x) { return x; }
static inline uint64_t ntohll(uint64_t x) { return x; }
#else
#error __BYTE_ORDER is neither __LITTLE_ENDIAN nor __BIG_ENDIAN
#endif

/* structure of test parameters */
struct config_t
{
    const char *dev_name; /* IB device name */
    char *server_name;    /* server host name */
    u_int32_t tcp_port;   /* server TCP port */
    int ib_port;          /* local IB port to work with */
    int gid_idx;          /* gid index to use */
};

/* structure to exchange data which is needed to connect the QPs */
struct cm_con_data_t
{
    uint64_t addr;   /* Buffer address */
    uint32_t rkey;   /* Remote key */
    uint32_t qp_num; /* QP number */
    uint16_t lid;    /* LID of the IB port */
    uint8_t gid[16]; /* gid */
} __attribute__((packed));

struct resources
{
    struct ibv_device_attr device_attr;         /* 设备属性 */
    struct ibv_port_attr port_attr;             /* IB端口属性 */
    struct cm_con_data_t remote_props[MAX_QPS]; /* 远程属性，每个QP对应一个 */
    struct ibv_context *ib_ctx;                 /* 设备上下文 */
    struct ibv_pd *pd;                          /* 保护域 */
    struct ibv_cq *cq[MAX_QPS];                 /* 完成队列，每个QP对应一个 */
    struct ibv_qp *qp[MAX_QPS];                 /* 队列对 */
    struct ibv_mr *mr[MAX_QPS];                 /* 内存注册，每个QP对应一个 */
    char *buf[MAX_QPS];                         /* 内存缓冲区指针，每个QP对应一个 */
    int sock;                                   /* TCP套接字文件描述符 */
};

struct config_t config = {
    NULL,  /* dev_name */
    NULL,  /* server_name */
    19875, /* tcp_port */
    1,     /* ib_port */
    -1     /* gid_idx */
};

static int sock_connect(const char *servername, int port)
{
    struct addrinfo *resolved_addr = NULL;
    struct addrinfo *iterator;
    char service[6];
    int sockfd = -1;
    int listenfd = 0;
    int tmp;
    struct addrinfo hints =
        {
            .ai_flags = AI_PASSIVE,
            .ai_family = AF_INET,
            .ai_socktype = SOCK_STREAM};
    if (sprintf(service, "%d", port) < 0)
        goto sock_connect_exit;
    /* Resolve DNS address, use sockfd as temp storage */
    sockfd = getaddrinfo(servername, service, &hints, &resolved_addr);
    if (sockfd < 0)
    {
        fprintf(stderr, "%s for %s:%d\n", gai_strerror(sockfd), servername, port);
        goto sock_connect_exit;
    }
    /* Search through results and find the one we want */
    for (iterator = resolved_addr; iterator; iterator = iterator->ai_next)
    {
        sockfd = socket(iterator->ai_family, iterator->ai_socktype, iterator->ai_protocol);
        if (sockfd >= 0)
        {
            if (servername)
            {
                /* Client mode. Initiate connection to remote */
                if ((tmp = connect(sockfd, iterator->ai_addr, iterator->ai_addrlen)))
                {
                    fprintf(stdout, "failed connect \n");
                    close(sockfd);
                    sockfd = -1;
                }
            }
            else
            {
                /* Server mode. Set up listening socket and accept a connection */
                listenfd = sockfd;
                sockfd = -1;
                if (bind(listenfd, iterator->ai_addr, iterator->ai_addrlen))
                    goto sock_connect_exit;
                listen(listenfd, 1);
                sockfd = accept(listenfd, NULL, 0);
            }
        }
    }
sock_connect_exit:
    if (listenfd)
        close(listenfd);
    if (resolved_addr)
        freeaddrinfo(resolved_addr);
    if (sockfd < 0)
    {
        if (servername)
            fprintf(stderr, "Couldn't connect to %s:%d\n", servername, port);
        else
        {
            perror("server accept");
            fprintf(stderr, "accept() failed\n");
        }
    }
    return sockfd;
}

int sock_sync_data(int sock, int xfer_size, char *local_data, char *remote_data)
{
    int rc;
    int read_bytes = 0;
    int total_read_bytes = 0;
    rc = write(sock, local_data, xfer_size);
    if (rc < xfer_size)
        fprintf(stderr, "Failed writing data during sock_sync_data\n");
    else
        rc = 0;
    while (!rc && total_read_bytes < xfer_size)
    {
        read_bytes = read(sock, remote_data, xfer_size);
        if (read_bytes > 0)
            total_read_bytes += read_bytes;
        else
            rc = read_bytes;
    }
    return rc;
}

static int poll_completion(struct resources *res, int qp_index)
{
    struct ibv_wc wc;
    unsigned long start_time_msec;
    unsigned long cur_time_msec;
    struct timeval cur_time;
    int poll_result;
    int rc = 0;
    /* poll the completion for a while before giving up of doing it .. */
    gettimeofday(&cur_time, NULL);
    start_time_msec = (cur_time.tv_sec * 1000) + (cur_time.tv_usec / 1000);
    do
    {
        poll_result = ibv_poll_cq(res->cq[qp_index], 1, &wc);
        gettimeofday(&cur_time, NULL);
        cur_time_msec = (cur_time.tv_sec * 1000) + (cur_time.tv_usec / 1000);
    } while ((poll_result == 0) && ((cur_time_msec - start_time_msec) < MAX_POLL_CQ_TIMEOUT));
    if (poll_result < 0)
    {
        /* poll CQ failed */
        fprintf(stderr, "poll CQ failed\n");
        rc = 1;
    }
    else if (poll_result == 0)
    { /* the CQ is empty */
        fprintf(stderr, "completion wasn't found in the CQ after timeout\n");
        rc = 1;
    }
    else
    {
        /* CQE found */
        fprintf(stdout, "completion was found in CQ with status 0x%x\n", wc.status);
        /* check the completion status (here we don't care about the completion opcode */
        if (wc.status != IBV_WC_SUCCESS)
        {
            fprintf(stderr, "got bad completion with status: 0x%x, vendor syndrome: 0x%x\n", wc.status,
                    wc.vendor_err);
            rc = 1;
        }
    }
    return rc;
}

static int post_send(struct resources *res, int opcode, int qp_index, size_t size)
{
    struct ibv_send_wr sr;
    struct ibv_sge sge;
    struct ibv_send_wr *bad_wr = NULL;
    int rc;

    /* 准备scatter/gather entry */
    memset(&sge, 0, sizeof(sge));
    sge.addr = (uintptr_t)res->buf[qp_index];
    sge.length = size;
    sge.lkey = res->mr[qp_index]->lkey;

    /* 准备发送工作请求 */
    memset(&sr, 0, sizeof(sr));
    sr.next = NULL;
    sr.wr_id = 0;
    sr.sg_list = &sge;
    sr.num_sge = 1;
    sr.opcode = opcode;
    sr.send_flags = IBV_SEND_SIGNALED;

    if (opcode != IBV_WR_SEND)
    {
        sr.wr.rdma.remote_addr = res->remote_props[qp_index].addr;
        sr.wr.rdma.rkey = res->remote_props[qp_index].rkey;
    }

    /* 使用指定的QP发布发送请求 */
    rc = ibv_post_send(res->qp[qp_index], &sr, &bad_wr);
    if (rc)
        fprintf(stderr, "failed to post SR for QP %d\n", qp_index);
    else
    {
        switch (opcode)
        {
        case IBV_WR_SEND:
            fprintf(stdout, "Send Request was posted for QP %d\n", qp_index);
            break;
        case IBV_WR_RDMA_READ:
            fprintf(stdout, "RDMA Read Request was posted for QP %d\n", qp_index);
            break;
        case IBV_WR_RDMA_WRITE:
            fprintf(stdout, "RDMA Write Request was posted for QP %d\n", qp_index);
            break;
        default:
            fprintf(stdout, "Unknown Request was posted for QP %d\n", qp_index);
            break;
        }
    }
    return rc;
}

static int post_receive(struct resources *res, int qp_index, size_t size)
{
    struct ibv_recv_wr rr;
    struct ibv_sge sge;
    struct ibv_recv_wr *bad_wr;
    int rc;

    /* 准备scatter/gather entry */
    memset(&sge, 0, sizeof(sge));
    sge.addr = (uintptr_t)res->buf[qp_index];
    sge.length = size;
    sge.lkey = res->mr[qp_index]->lkey;

    /* 准备接收工作请求 */
    memset(&rr, 0, sizeof(rr));
    rr.next = NULL;
    rr.wr_id = 0;
    rr.sg_list = &sge;
    rr.num_sge = 1;

    /* 使用指定的QP发布接收请求 */
    rc = ibv_post_recv(res->qp[qp_index], &rr, &bad_wr);
    if (rc)
        fprintf(stderr, "failed to post RR for QP %d\n", qp_index);
    else
        fprintf(stdout, "Receive Request was posted for QP %d\n", qp_index);
    return rc;
}

static void resources_init(struct resources *res)
{
    memset(res, 0, sizeof *res);
    res->sock = -1;
}

static int resources_create(struct resources *res, int num_qps)
{
    struct ibv_device **dev_list = NULL;
    struct ibv_qp_init_attr qp_init_attr;
    struct ibv_device *ib_dev = NULL;
    size_t size;
    int i;
    int mr_flags = 0;
    int cq_size = 0;
    int num_devices;
    int rc = 0;

    /* 建立TCP连接 */
    if (config.server_name)
    {
        res->sock = sock_connect(config.server_name, config.tcp_port);
        if (res->sock < 0)
        {
            fprintf(stderr, "failed to establish TCP connection to server %s, port %d\n",
                    config.server_name, config.tcp_port);
            rc = -1;
            goto resources_create_exit;
        }
    }
    else
    {
        fprintf(stdout, "waiting on port %d for TCP connection\n", config.tcp_port);
        res->sock = sock_connect(NULL, config.tcp_port);
        if (res->sock < 0)
        {
            fprintf(stderr, "failed to establish TCP connection with client on port %d\n",
                    config.tcp_port);
            rc = -1;
            goto resources_create_exit;
        }
    }
    fprintf(stdout, "TCP connection was established\n");
    fprintf(stdout, "searching for IB devices in host\n");
    /* 获取设备列表 */
    dev_list = ibv_get_device_list(&num_devices);
    if (!dev_list)
    {
        fprintf(stderr, "failed to get IB devices list\n");
        rc = 1;
        goto resources_create_exit;
    }
    if (!num_devices)
    {
        fprintf(stderr, "found %d device(s)\n", num_devices);
        rc = 1;
        goto resources_create_exit;
    }
    fprintf(stdout, "found %d device(s)\n", num_devices);

    /* 选择设备 */
    for (i = 0; i < num_devices; i++)
    {
        if (!config.dev_name)
        {
            config.dev_name = strdup(ibv_get_device_name(dev_list[i]));
            fprintf(stdout, "device not specified, using first one found: %s\n", config.dev_name);
        }
        if (!strcmp(ibv_get_device_name(dev_list[i]), config.dev_name))
        {
            ib_dev = dev_list[i];
            break;
        }
    }
    if (!ib_dev)
    {
        fprintf(stderr, "IB device %s wasn't found\n", config.dev_name);
        rc = 1;
        goto resources_create_exit;
    }

    /* 打开设备上下文 */
    res->ib_ctx = ibv_open_device(ib_dev);
    if (!res->ib_ctx)
    {
        fprintf(stderr, "failed to open device %s\n", config.dev_name);
        rc = 1;
        goto resources_create_exit;
    }

    /* 释放设备列表 */
    ibv_free_device_list(dev_list);
    dev_list = NULL;
    ib_dev = NULL;

    /* 查询端口属性 */
    if (ibv_query_port(res->ib_ctx, config.ib_port, &res->port_attr))
    {
        fprintf(stderr, "ibv_query_port on port %u failed\n", config.ib_port);
        rc = 1;
        goto resources_create_exit;
    }

    /* 分配保护域 */
    res->pd = ibv_alloc_pd(res->ib_ctx);
    if (!res->pd)
    {
        fprintf(stderr, "ibv_alloc_pd failed\n");
        rc = 1;
        goto resources_create_exit;
    }

    mr_flags = IBV_ACCESS_LOCAL_WRITE | IBV_ACCESS_REMOTE_READ | IBV_ACCESS_REMOTE_WRITE;

    for (i = 0; i < 1999; i++)
    {
        /* 创建完成队列 */
        cq_size = 1; // 每个QP的完成队列大小
        res->cq[i] = ibv_create_cq(res->ib_ctx, cq_size, NULL, NULL, 0);
        if (!res->cq[i])
        {
            fprintf(stderr, "failed to create CQ with %u entries for QP %d\n", cq_size, i);
            rc = 1;
            goto resources_create_exit;
        }

        /* 为每个QP分配内存缓冲区 */
        res->buf[i] = (char *)malloc(SMALL_MSG_SIZE);
        if (!res->buf[i])
        {
            fprintf(stderr, "failed to malloc %Zu bytes to memory buffer for QP %d\n", SMALL_MSG_SIZE, i);
            rc = 1;
            goto resources_create_exit;
        }
        memset(res->buf[i], 0, SMALL_MSG_SIZE);

        /* 在服务器端放置初始消息 */
        if (!config.server_name)
        {
            strcpy(res->buf[i], MSG);
            fprintf(stdout, "QP %d: going to send the message: '%s'\n", i, res->buf[i]);
        }
        else
        {
            memset(res->buf[i], 0, SMALL_MSG_SIZE);
        }

        /* 注册内存缓冲区 */
        res->mr[i] = ibv_reg_mr(res->pd, res->buf[i], SMALL_MSG_SIZE, mr_flags);
        if (!res->mr[i])
        {
            fprintf(stderr, "ibv_reg_mr failed with mr_flags=0x%x for QP %d\n", mr_flags, i);
            rc = 1;
            goto resources_create_exit;
        }
        fprintf(stdout, "MR for QP %d was registered with addr=%p, lkey=0x%x, rkey=0x%x, flags=0x%x\n",
                i, res->buf[i], res->mr[i]->lkey, res->mr[i]->rkey, mr_flags);

        /* 创建队列对 */
        memset(&qp_init_attr, 0, sizeof(qp_init_attr));
        qp_init_attr.qp_type = IBV_QPT_RC;
        qp_init_attr.sq_sig_all = 1;
        qp_init_attr.send_cq = res->cq[i];
        qp_init_attr.recv_cq = res->cq[i];
        qp_init_attr.cap.max_send_wr = 1;
        qp_init_attr.cap.max_recv_wr = 1;
        qp_init_attr.cap.max_send_sge = 1;
        qp_init_attr.cap.max_recv_sge = 1;

        res->qp[i] = ibv_create_qp(res->pd, &qp_init_attr);
        if (!res->qp[i])
        {
            fprintf(stderr, "failed to create QP %d\n", i);
            rc = 1;
            goto resources_create_exit;
        }
        fprintf(stdout, "QP %d was created, QP number=0x%x\n", i, res->qp[i]->qp_num);
    }

    for (i = 1999; i < MAX_QPS; i++)
    {
        /* 创建完成队列 */
        cq_size = (i == 1999) ? 2 : 1;
        res->cq[i] = ibv_create_cq(res->ib_ctx, cq_size, NULL, NULL, 0);
        if (!res->cq[i])
        {
            fprintf(stderr, "failed to create CQ with %u entries for QP %d\n", cq_size, i);
            rc = 1;
            goto resources_create_exit;
        }

        /* 为每个QP分配内存缓冲区 */
        res->buf[i] = (char *)malloc(LARGE_MSG_SIZE);
        if (!res->buf[i])
        {
            fprintf(stderr, "failed to malloc %Zu bytes to memory buffer for QP %d\n", LARGE_MSG_SIZE, i);
            rc = 1;
            goto resources_create_exit;
        }
        memset(res->buf[i], 0, LARGE_MSG_SIZE);

        /* 在服务器端放置初始消息 */
        if (!config.server_name)
        {
            strcpy(res->buf[i], MSG);
            fprintf(stdout, "QP %d: going to send the large message\n", i);
        }
        else
        {
            memset(res->buf[i], 0, LARGE_MSG_SIZE);
        }

        /* 注册大内存缓冲区 */
        res->mr[i] = ibv_reg_mr(res->pd, res->buf[i], LARGE_MSG_SIZE, mr_flags);
        if (!res->mr[i])
        {
            fprintf(stderr, "ibv_reg_mr failed with mr_flags=0x%x for QP %d\n", mr_flags, i);
            rc = 1;
            goto resources_create_exit;
        }
        fprintf(stdout, "MR for QP %d was registered with addr=%p, lkey=0x%x, rkey=0x%x, flags=0x%x\n",
                i, res->buf[i], res->mr[i]->lkey, res->mr[i]->rkey, mr_flags);

        /* 创建队列对 */
        memset(&qp_init_attr, 0, sizeof(qp_init_attr));
        qp_init_attr.qp_type = IBV_QPT_RC;
        qp_init_attr.sq_sig_all = 1;
        qp_init_attr.send_cq = res->cq[i];
        qp_init_attr.recv_cq = res->cq[i];
        qp_init_attr.cap.max_send_wr = 1;
        qp_init_attr.cap.max_recv_wr = 1;
        qp_init_attr.cap.max_send_sge = 1;
        qp_init_attr.cap.max_recv_sge = 1;

        res->qp[i] = ibv_create_qp(res->pd, &qp_init_attr);
        if (!res->qp[i])
        {
            fprintf(stderr, "failed to create QP %d\n", i);
            rc = 1;
            goto resources_create_exit;
        }
        fprintf(stdout, "QP %d was created, QP number=0x%x\n", i, res->qp[i]->qp_num);
    }

resources_create_exit:
    if (rc)
    {
        /* 释放已分配的资源 */
        for (i = 0; i < num_qps; i++)
        {
            if (res->qp[i])
            {
                ibv_destroy_qp(res->qp[i]);
                res->qp[i] = NULL;
            }
            if (res->mr[i])
            {
                ibv_dereg_mr(res->mr[i]);
                res->mr[i] = NULL;
            }
            if (res->buf[i])
            {
                free(res->buf[i]);
                res->buf[i] = NULL;
            }
            if (res->cq[i])
            {
                ibv_destroy_cq(res->cq[i]);
                res->cq[i] = NULL;
            }
        }
        if (res->pd)
        {
            ibv_dealloc_pd(res->pd);
            res->pd = NULL;
        }
        if (res->ib_ctx)
        {
            ibv_close_device(res->ib_ctx);
            res->ib_ctx = NULL;
        }
        if (dev_list)
        {
            ibv_free_device_list(dev_list);
            dev_list = NULL;
        }
        if (res->sock >= 0)
        {
            if (close(res->sock))
                fprintf(stderr, "failed to close socket\n");
            res->sock = -1;
        }
    }
    return rc;
}

static int modify_qp_to_init(struct ibv_qp *qp)
{
    struct ibv_qp_attr attr;
    int flags;
    int rc;
    memset(&attr, 0, sizeof(attr));
    attr.qp_state = IBV_QPS_INIT;
    attr.port_num = config.ib_port;
    attr.pkey_index = 0;
    attr.qp_access_flags = IBV_ACCESS_LOCAL_WRITE | IBV_ACCESS_REMOTE_READ | IBV_ACCESS_REMOTE_WRITE;
    flags = IBV_QP_STATE | IBV_QP_PKEY_INDEX | IBV_QP_PORT | IBV_QP_ACCESS_FLAGS;
    rc = ibv_modify_qp(qp, &attr, flags);
    if (rc)
        fprintf(stderr, "failed to modify QP state to INIT\n");
    return rc;
}

static int modify_qp_to_rtr(struct ibv_qp *qp, uint32_t remote_qpn, uint16_t dlid, uint8_t *dgid)
{
    struct ibv_qp_attr attr;
    int flags;
    int rc;
    memset(&attr, 0, sizeof(attr));
    attr.qp_state = IBV_QPS_RTR;
    attr.path_mtu = IBV_MTU_256;
    attr.dest_qp_num = remote_qpn;
    attr.rq_psn = 0;
    attr.max_dest_rd_atomic = 1;
    attr.min_rnr_timer = 0x12;
    attr.ah_attr.is_global = 0;
    attr.ah_attr.dlid = dlid;
    attr.ah_attr.sl = 0;
    attr.ah_attr.src_path_bits = 0;
    attr.ah_attr.port_num = config.ib_port;
    if (config.gid_idx >= 0)
    {
        attr.ah_attr.is_global = 1;
        attr.ah_attr.port_num = 1;
        memcpy(&attr.ah_attr.grh.dgid, dgid, 16);
        attr.ah_attr.grh.flow_label = 0;
        attr.ah_attr.grh.hop_limit = 1;
        attr.ah_attr.grh.sgid_index = config.gid_idx;
        attr.ah_attr.grh.traffic_class = 0;
    }
    flags = IBV_QP_STATE | IBV_QP_AV | IBV_QP_PATH_MTU | IBV_QP_DEST_QPN |
            IBV_QP_RQ_PSN | IBV_QP_MAX_DEST_RD_ATOMIC | IBV_QP_MIN_RNR_TIMER;
    rc = ibv_modify_qp(qp, &attr, flags);
    if (rc)
        fprintf(stderr, "failed to modify QP state to RTR\n");
    return rc;
}

static int modify_qp_to_rts(struct ibv_qp *qp)
{
    struct ibv_qp_attr attr;
    int flags;
    int rc;
    memset(&attr, 0, sizeof(attr));
    attr.qp_state = IBV_QPS_RTS;
    attr.timeout = 0x12;
    attr.retry_cnt = 6;
    attr.rnr_retry = 0;
    attr.sq_psn = 0;
    attr.max_rd_atomic = 1;
    flags = IBV_QP_STATE | IBV_QP_TIMEOUT | IBV_QP_RETRY_CNT |
            IBV_QP_RNR_RETRY | IBV_QP_SQ_PSN | IBV_QP_MAX_QP_RD_ATOMIC;
    rc = ibv_modify_qp(qp, &attr, flags);
    if (rc)
        fprintf(stderr, "failed to modify QP state to RTS\n");
    return rc;
}

static int connect_qp(struct resources *res, int qp_index)
{
    struct cm_con_data_t local_con_data;
    struct cm_con_data_t remote_con_data;
    struct cm_con_data_t tmp_con_data;
    int rc = 0;
    char temp_char;
    union ibv_gid my_gid;
    if (config.gid_idx >= 0)
    {
        rc = ibv_query_gid(res->ib_ctx, config.ib_port, config.gid_idx, &my_gid);
        if (rc)
        {
            fprintf(stderr, "could not get gid for port %d, index %d\n", config.ib_port, config.gid_idx);
            return rc;
        }
    }
    else
        memset(&my_gid, 0, sizeof my_gid);

    /* 使用TCP套接字交换连接信息 */
    local_con_data.addr = htonll((uintptr_t)res->buf[qp_index]);
    local_con_data.rkey = htonl(res->mr[qp_index]->rkey);
    local_con_data.qp_num = htonl(res->qp[qp_index]->qp_num);
    local_con_data.lid = htons(res->port_attr.lid);
    memcpy(local_con_data.gid, &my_gid, 16);
    fprintf(stdout, "\nLocal LID for QP %d = 0x%x\n", qp_index, res->port_attr.lid);

    if (sock_sync_data(res->sock, sizeof(struct cm_con_data_t), (char *)&local_con_data, (char *)&tmp_con_data) < 0)
    {
        fprintf(stderr, "failed to exchange connection data between sides for QP %d\n", qp_index);
        rc = 1;
        goto connect_qp_exit;
    }

    remote_con_data.addr = ntohll(tmp_con_data.addr);
    remote_con_data.rkey = ntohl(tmp_con_data.rkey);
    remote_con_data.qp_num = ntohl(tmp_con_data.qp_num);
    remote_con_data.lid = ntohs(tmp_con_data.lid);
    memcpy(remote_con_data.gid, tmp_con_data.gid, 16);

    /* 保存远程属性 */
    res->remote_props[qp_index] = remote_con_data;
    fprintf(stdout, "Remote address for QP %d = 0x%" PRIx64 "\n", qp_index, remote_con_data.addr);
    fprintf(stdout, "Remote rkey for QP %d = 0x%x\n", qp_index, remote_con_data.rkey);
    fprintf(stdout, "Remote QP number for QP %d = 0x%x\n", qp_index, remote_con_data.qp_num);
    fprintf(stdout, "Remote LID for QP %d = 0x%x\n", qp_index, remote_con_data.lid);

    if (config.gid_idx >= 0)
    {
        uint8_t *p = remote_con_data.gid;
        fprintf(stdout, "Remote GID for QP %d = %02x:%02x:%02x:%02x:%02x:%02x:%02x:%02x:%02x:%02x:%02x:%02x:%02x:%02x:%02x:%02x\n ",
                qp_index, p[0], p[1], p[2], p[3], p[4], p[5], p[6], p[7], p[8], p[9], p[10], p[11], p[12], p[13], p[14], p[15]);
    }

    /* 修改QP状态到INIT */
    rc = modify_qp_to_init(res->qp[qp_index]);
    if (rc)
    {
        fprintf(stderr, "change QP %d state to INIT failed\n", qp_index);
        goto connect_qp_exit;
    }


    /* 修改QP状态到RTR */
    rc = modify_qp_to_rtr(res->qp[qp_index], remote_con_data.qp_num, remote_con_data.lid, remote_con_data.gid);
    if (rc)
    {
        fprintf(stderr, "failed to modify QP %d state to RTR\n", qp_index);
        goto connect_qp_exit;
    }

    /* 修改QP状态到RTS */
    rc = modify_qp_to_rts(res->qp[qp_index]);
    if (rc)
    {
        fprintf(stderr, "failed to modify QP %d state to RTS\n", qp_index);
        goto connect_qp_exit;
    }
    fprintf(stdout, "QP %d state was changed to RTS\n", qp_index);

    /* 同步状态 */
    if (sock_sync_data(res->sock, 1, "Q", &temp_char))
    {
        fprintf(stderr, "sync error after QPs were moved to RTS for QP %d\n", qp_index);
        rc = 1;
    }

connect_qp_exit:
    return rc;
}

static int resources_destroy(struct resources *res, int num_qps)
{
    int rc = 0;

    for (int i = 0; i < num_qps; i++)
    {
        if (res->qp[i])
        {
            if (ibv_destroy_qp(res->qp[i]))
            {
                fprintf(stderr, "failed to destroy QP %d\n", i);
                rc = 1;
            }
        }
        if (res->mr[i])
        {
            if (ibv_dereg_mr(res->mr[i]))
            {
                fprintf(stderr, "failed to deregister MR for QP %d\n", i);
                rc = 1;
            }
        }
        if (res->buf[i])
        {
            free(res->buf[i]);
            res->buf[i] = NULL;
        }
        if (res->cq[i])
        {
            if (ibv_destroy_cq(res->cq[i]))
            {
                fprintf(stderr, "failed to destroy CQ for QP %d\n", i);
                rc = 1;
            }
        }
    }
    if (res->pd)
    {
        if (ibv_dealloc_pd(res->pd))
        {
            fprintf(stderr, "failed to deallocate PD\n");
            rc = 1;
        }
    }
    if (res->ib_ctx)
    {
        if (ibv_close_device(res->ib_ctx))
        {
            fprintf(stderr, "failed to close device context\n");
            rc = 1;
        }
    }
    if (res->sock >= 0)
    {
        if (close(res->sock))
        {
            fprintf(stderr, "failed to close socket\n");
            rc = 1;
        }
    }
    return rc;
}

static void print_config(void)
{
    fprintf(stdout, " ------------------------------------------------\n");
    fprintf(stdout, " Device name : \"%s\"\n", config.dev_name);
    fprintf(stdout, " IB port : %u\n", config.ib_port);
    if (config.server_name)
        fprintf(stdout, " IP : %s\n", config.server_name);
    fprintf(stdout, " TCP port : %u\n", config.tcp_port);
    if (config.gid_idx >= 0)
        fprintf(stdout, " GID index : %u\n", config.gid_idx);
    fprintf(stdout, " ------------------------------------------------\n\n");
}

static void usage(const char *argv0)
{
    fprintf(stdout, "Usage:\n");
    fprintf(stdout, " %s start a server and wait for connection\n", argv0);
    fprintf(stdout, " %s <host> connect to server at <host>\n", argv0);
    fprintf(stdout, "\n");
    fprintf(stdout, "Options:\n");
    fprintf(stdout, " -p, --port <port> listen on/connect to port <port> (default 18515)\n");
    fprintf(stdout, " -d, --ib-dev <dev> use IB device <dev> (default first device found)\n");
    fprintf(stdout, " -i, --ib-port <port> use port <port> of IB device (default 1)\n");
    fprintf(stdout, " -g, --gid_idx <git index> gid index to be used in GRH (default not used)\n");
}

int main(int argc, char *argv[])
{
    struct resources res;
    int rc = 1;
    char temp_char;
    int num_qps = 2001; // QP数量

    /* 解析命令行参数 */
    while (1)
    {
        int c;
        static struct option long_options[] = {
            {.name = "port", .has_arg = 1, .val = 'p'},
            {.name = "ib-dev", .has_arg = 1, .val = 'd'},
            {.name = "ib-port", .has_arg = 1, .val = 'i'},
            {.name = "gid-idx", .has_arg = 1, .val = 'g'},
            {.name = NULL, .has_arg = 0, .val = '\0'}};
        c = getopt_long(argc, argv, "p:d:i:g:", long_options, NULL);
        if (c == -1)
            break;
        switch (c)
        {
        case 'p':
            config.tcp_port = strtoul(optarg, NULL, 0);
            break;
        case 'd':
            config.dev_name = strdup(optarg);
            break;
        case 'i':
            config.ib_port = strtoul(optarg, NULL, 0);
            if (config.ib_port < 0)
            {
                usage(argv[0]);
                return 1;
            }
            break;
        case 'g':
            config.gid_idx = strtoul(optarg, NULL, 0);
            if (config.gid_idx < 0)
            {
                usage(argv[0]);
                return 1;
            }
            break;
        default:
            usage(argv[0]);
            return 1;
        }
    }

    /* 解析服务器名参数 */
    if (optind == argc - 1)
        config.server_name = argv[optind];
    if (config.server_name)
    {
        printf("servername=%s\n", config.server_name);
    }
    else if (optind < argc)
    {
        usage(argv[0]);
        return 1;
    }

    /* 打印配置信息 */
    print_config();

    /* 初始化资源 */
    resources_init(&res);

    /* 创建资源 */
    if (resources_create(&res, num_qps))
    {
        fprintf(stderr, "failed to create resources\n");
        goto main_exit;
    }

    /* 连接所有QPs */
    for (int i = 0; i < num_qps; i++)
    {
        if (connect_qp(&res, i))
        {
            fprintf(stderr, "failed to connect QP %d\n", i);
            goto main_exit;
        }
    }

    for (int i = 0; i < 2000; i++)
    {
        if (!config.server_name)
        {
            if (post_send(&res, IBV_WR_RDMA_WRITE, i, SMALL_MSG_SIZE))
            {
                fprintf(stderr, "failed to post rdma write request for QP %d\n", i);
                goto main_exit;
            }
        }

        if (!config.server_name)
        {
            if (poll_completion(&res, i))
            {
                fprintf(stderr, "poll completion failed for QP %d\n", i);
                goto main_exit;
            }
        }
    }

    /* 同步状态，确保服务器和客户端都准备好进行下一个测试 */
    if (sock_sync_data(res.sock, 1, "R", &temp_char))
    {
        fprintf(stderr, "sync error before RDMA ops\n");
        rc = 1;
        goto main_exit;
    }

    struct timeval start_time, end_time, start_time_2, end_time_2;
    gettimeofday(&start_time, NULL);

    if (!config.server_name)
    {
        if (post_send(&res, IBV_WR_RDMA_WRITE, 1999, MID_MSG_SIZE))
        {
            fprintf(stderr, "failed to post large rdma write request for QP 1999\n");
            goto main_exit;
        }
    }

    if (!config.server_name)
    {
        if (poll_completion(&res, 1999))
        {
            fprintf(stderr, "poll completion failed for QP 1999\n");
            goto main_exit;
        }
    }

    gettimeofday(&end_time, NULL);
    long time_diff_500th = (end_time.tv_sec - start_time.tv_sec) * 1000000 + (end_time.tv_usec - start_time.tv_usec);
    fprintf(stdout, "Time taken for 2000th QP : %ld microseconds\n", time_diff_500th);

    

    gettimeofday(&start_time_2, NULL);

    if (!config.server_name)
    {
        if (post_send(&res, IBV_WR_RDMA_WRITE, 2000, MID_MSG_SIZE))
        {
            fprintf(stderr, "failed to post large rdma write request for QP 2000\n");
            goto main_exit;
        }
    }

    if (!config.server_name)
    {
        if (poll_completion(&res, 2000))
        {
            fprintf(stderr, "poll completion failed for QP 2000\n");
            goto main_exit;
        }
    }

    gettimeofday(&end_time_2, NULL);
    long time_diff_501st = (end_time_2.tv_sec - start_time_2.tv_sec) * 1000000 + (end_time_2.tv_usec - start_time_2.tv_usec);
    fprintf(stdout, "Time taken for 2001st QP : %ld microseconds\n", time_diff_501st);

    rc = 0;

main_exit:
    if (resources_destroy(&res, num_qps))
    {
        fprintf(stderr, "failed to destroy resources\n");
        rc = 1;
    }
    if (config.dev_name)
        free((char *)config.dev_name);
    fprintf(stdout, "\ntest result is %d\n", rc);
    return rc;
}
